import torchsummary as summary 
import torchvision
import datetime
import torch
from PIL import Image
import os

'''定义公共层：conv+bn+relu'''
def CBR(in_channels,out_channels):
    cbr=torch.nn.Sequential(
        torch.nn.Conv2d(
            in_channels, 
            out_channels,
            kernel_size=3,
            stride=1,
            padding=1
        ),
        torch.nn.BatchNorm2d(out_channels),
        torch.nn.ReLU(inplace=True)
    )
    return cbr

'''定义VGG'''
class VGG(torch.nn.Module):
    
    def __init__(self, block_nums,classes_num):
        
        super(VGG,self).__init__()
        
        self.block1=self._make_layers(
            in_channels=3,out_channels=64,block_num=block_nums[0]
            )
        self.block2=self._make_layers(
            in_channels=64,out_channels=128,block_num=block_nums[1]
            )
        self.block3=self._make_layers(
            in_channels=128,out_channels=256,block_num=block_nums[2]
            )
        self.block4=self._make_layers(
            in_channels=256,out_channels=512,block_num=block_nums[3]
            )
        self.block5=self._make_layers(
            in_channels=512,out_channels=512,block_num=block_nums[4]
            )
        
        self.classifier=torch.nn.Sequential(
            
            torch.nn.Linear(512*7*7, 4096),
            torch.nn.ReLU(inplace=True),
            torch.nn.Dropout(0.5),
            
            torch.nn.Linear(4096, 4096),
            torch.nn.ReLU(inplace=True),
            torch.nn.Dropout(0.5),
            
            torch.nn.Linear(4096, classes_num)
        )
    
    def _make_layers(self,in_channels,out_channels,block_num):
        
        blocks=[]
        blocks.append(CBR(in_channels,out_channels))
        
        for i in range(1,block_num):
            
            blocks.append(CBR(out_channels,out_channels))
        
        blocks.append(torch.nn.MaxPool2d(kernel_size=2,stride=2))
            
        return torch.nn.Sequential(*blocks)
        
    def forward(self,x):
        
        x=self.block1(x)
        x=self.block2(x)
        x=self.block3(x)
        x=self.block4(x)
        x=self.block5(x)
        x=torch.flatten(x,start_dim=1)
        out=self.classifier(x)
        
        return out

# start train and val model
def train(
        model_params_path,
        train_dataset_dir,
        val_dataset_dir=None,
        batch_size=64,
        epoch_total=10,
        classes_num=10,
        learning_rate=1e-3,
        block_nums=[2,2,3,3,3]
    ):
    # preconditioning dataset 
    transform = torchvision.transforms.Compose([
        torchvision.transforms.RandomResizedCrop(224),
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
    ])

    train_dataset=torchvision.datasets.ImageFolder(
        root=train_dataset_dir,transform=transform
    )
    
    # set dataset loader
    train_loader=torch.utils.data.DataLoader(
        dataset=train_dataset,batch_size=batch_size,shuffle=True
    )

    # set model
    model=VGG(block_nums=block_nums,classes_num=classes_num)
    print(model)
    summary.summary(model, input_size=(3,224,224),device="cpu")

    # set loss and criterion
    criterion = torch.nn.CrossEntropyLoss()
    optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)

    # start train
    step_total=len(train_loader)
    for epoch in range(epoch_total):
        for step,data in enumerate(train_loader,start=0):
            image,label=data
            pred=model(image)
            loss=criterion(pred,label)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            print(f"{datetime.datetime.now()}>")
            print(f"Epoch:[{epoch+1}/{epoch_total}],",end="")
            print(f"Step:[{step+1}/{step_total}],",end="")
            print(f"loss:{round(loss.item(),4)}")
    torch.save(model.state_dict(), model_params_path)

    # val model
    if val_dataset_dir is not None:
        val(val_dataset_dir,model,model_params_path,batch_size=batch_size)
    return

# start val model
def val(
        dataset_dir, 
        model, 
        model_params_path,
        batch_size=64
    ):
    # dataset loader
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((224,224)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
    ])
    val_dataset=torchvision.datasets.ImageFolder(
        root=dataset_dir,transform=transform
    )
    val_loader=torch.utils.data.DataLoader(
        dataset=val_dataset,batch_size=batch_size,shuffle=True
    )

    dataiter = iter(val_loader)
    images, labels = dataiter.next()

    model.load_state_dict(torch.load(model_params_path))
    outputs = model(images)

    correct = 0
    total = 0
    with torch.no_grad():
        for data in val_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    with torch.no_grad():
        for data in val_loader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            c = (predicted == labels).squeeze()
            for i in range(4):
                label = labels[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1

    return

# start predict model
def predict(predict_loader,model_path):
    transforms=torchvision.transforms.Compose([
        torchvision.transforms.Resize((64,64)),
        torchvision.transforms.ToTensor()
    ])

    # model=VGG(block_nums=block_nums,classes_num=classes_num)
    model=torch.load(model_path,map_location=torch.device("cpu"))
    for name in os.listdir(predict_loader):
        path = os.path.join(name)
        image=Image.open(path)
        image=transforms(image)
        image=torch.reshape(image,(1,3,64,64))
    return

# start
def main():
    '''train and val model'''
    model_params_path=r"E:\Yang\train_model_20220423.pth"
    train_dataset_dir=r""
    val_dataset_dir=None
    classes_num=10
    train( 
        model_params_path,
        train_dataset_dir,
        val_dataset_dir,
        classes_num,
     )
    '''predict model'''

    model_path=r"E:\Yang\train_model_20220423.pth"
    predict()
    return


if __name__ == '__main__':
    main()

            

        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        